def clean_node_id(node_id):
    """去除节点ID中的空白字符并返回标准化后的字符串"""
    return node_id.strip()

def process_files(dm_file_path, dm_label_file_path, output_file_path):
    # 创建一个集合来保存所有出现在 DM.txt 文件中的节点的数字部分
    nodes_set = set()

    # 读取 DM.txt 文件并填充集合
    with open(dm_file_path, 'r') as dm_file:
        for line in dm_file:
            parts = line.strip().split()
            if len(parts) == 2:
                d_node, m_node = parts
                nodes_set.add(clean_node_id(d_node[1:]))  # 去掉前缀 D 并清洗
                nodes_set.add(clean_node_id(m_node[1:]))  # 去掉前缀 M 并清洗

    # 打印集合大小以确认
    print(f"Total unique nodes from DM.txt: {len(nodes_set)}")

    # 过滤 DM_label.txt 文件并写入输出文件
    matched_lines = 0
    unmatched_nodes = set()  # 记录未匹配到的节点ID
    invalid_lines = 0  # 记录无效行的数量

    with open(dm_label_file_path, 'r') as dm_label_file, open(output_file_path, 'w') as output_file:
        for line in dm_label_file:
            parts = line.strip().split('\t')
            if len(parts) == 2:
                label, vector = parts
                node_id = clean_node_id(label[1:])  # 去掉前缀 p 或 m 并清洗

                # 检查节点ID的数字部分是否在集合中
                if node_id in nodes_set:
                    output_file.write(line)
                    matched_lines += 1
                else:
                    unmatched_nodes.add(node_id)
            else:
                invalid_lines += 1

    print(f"Total lines written to {output_file_path}: {matched_lines}")
    print(f"Number of unmatched nodes: {len(unmatched_nodes)}")
    print(f"Number of invalid lines in {dm_label_file_path}: {invalid_lines}")

    # 如果需要，可以打印未匹配的节点ID进行进一步分析
    if unmatched_nodes:
        print("Unmatched nodes (first 10):", list(unmatched_nodes)[:10])

if __name__ == "__main__":
    dm_file_path = r'C:\Users\C\Desktop\NSEVis\src\assets\AM.txt'
    dm_label_file_path = r'C:\Users\C\Desktop\NSEVis\src\assets\AM_label.txt'
    output_file_path = 'filtered_AM_label.txt'

    process_files(dm_file_path, dm_label_file_path, output_file_path)

    print(f"Filtered labels have been saved to {output_file_path}")